import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision import datasets
import matplotlib.pyplot as plt
import numpy as np

from fedSNN import SNN, transform_cifar_gray, transform_mnist


def imshow(img):
    img = img / 2 + 0.5  # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


def evaluate_model(model, test_loader, device='cuda'):
    model.eval()
    correct = 0
    total = 0
    dataiter = iter(test_loader)
    images, labels = next(dataiter)
    images, labels = images.to(device), labels.to(device)

    with torch.no_grad():
        outputs = model(images.view(images.size(0), -1))
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f"Test Accuracy: {accuracy:.2f}%")

    # Show images and predictions
    imshow(torchvision.utils.make_grid(images.cpu()))
    print('GroundTruth: ', ' '.join('%5s' % labels[j].item() for j in range(4)))
    print('Predicted: ', ' '.join('%5s' % predicted[j].item() for j in range(4)))

    return accuracy


if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dataset = 'cifar10'

    if dataset == 'mnist':
        test_loader = DataLoader(datasets.MNIST(root='./data', train=False, download=True, transform=transform_mnist),
                                 batch_size=32, shuffle=False)
        input_size = 784
    elif dataset == 'fashion_mnist':
        test_loader = DataLoader(
            datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform_mnist), batch_size=32,
            shuffle=False)
        input_size = 784
    elif dataset == 'cifar10':
        test_loader = DataLoader(
            datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_cifar_gray), batch_size=32,
            shuffle=False)
        input_size = 1024
    elif dataset == 'cifar100':
        test_loader = DataLoader(
            datasets.CIFAR100(root='./data', train=False, download=True, transform=transform_cifar_gray), batch_size=32,
            shuffle=False)
        input_size = 1024
    else:
        raise ValueError(f"Unknown dataset: {dataset}")

    model = SNN(input_size=input_size, hidden_size1=512, hidden_size2=256, output_size=10).to(device)
    model.load_state_dict(torch.load('best_model.pth'))
    evaluate_model(model, test_loader, device)